use sq_global::{API_STOCK_COMP, API_STOCK_DAILY};

use polars::error::PolarsResult;
use polars::prelude::{
    Column, DataType, GetOutput, IntoColumn, IntoLazy, ListNameSpaceExtension, NULL, NamedFrom,
    RoundMode, Series, SortMultipleOptions, col, lit, when,
};
use sq_macro::test_with_ctx;
use sq_polars::{get_col_idx, get_str_col};
use tokio::try_join;

#[test_with_ctx]
pub async fn test() {
    run().await.expect("ssss");
}

pub async fn run() -> anyhow::Result<Vec<String>> {
    let (com_df, daily_dfs) = try_join!(
        sq_tushare::load_data(API_STOCK_COMP, Vec::new()),
        sq_tushare::recent_dates(API_STOCK_DAILY, 20)
    )?;
    let daily_df = sq_polars::concat_dfs(daily_dfs)?;

    let frame = daily_df
        .select([col("ts_code"), col("pct_chg"), col("trade_date")])
        .sort(
            ["trade_date"],
            SortMultipleOptions::new().with_order_descending(true),
        )
        .group_by_stable(["ts_code"])
        .agg([
            col("pct_chg"),
            col("pct_chg")
                .apply(calc_5d_10d, GetOutput::from_type(DataType::Float64))
                .alias("5d_10d_lines"),
        ])
        .filter(col("5d_10d_lines").list().len().eq(lit(4)))
        .with_columns([
            ((col("5d_10d_lines").list().get(lit(1), false)
                - col("5d_10d_lines").list().get(lit(3), false))
                / ((col("5d_10d_lines").list().get(lit(0), false)
                    - col("5d_10d_lines").list().get(lit(2), false))
                .abs()))
            .alias("trend_degree"),
            when(
                col("5d_10d_lines")
                    .list()
                    .get(lit("0"), false)
                    .gt(col("5d_10d_lines").list().get(lit("2"), false))
                    .and(
                        col("5d_10d_lines")
                            .list()
                            .get(lit("1"), false)
                            .lt(col("5d_10d_lines").list().get(lit("3"), false)),
                    ),
            )
            .then(lit("否"))
            .when(
                col("5d_10d_lines")
                    .list()
                    .get(lit("0"), false)
                    .lt(col("5d_10d_lines").list().get(lit("2"), false))
                    .and(
                        col("5d_10d_lines")
                            .list()
                            .get(lit("1"), false)
                            .gt(col("5d_10d_lines").list().get(lit("3"), false)),
                    ),
            )
            .then(lit("是"))
            .otherwise(lit(NULL).cast(DataType::String))
            .alias("up"),
        ])
        .filter(col("up").is_not_null())
        .inner_join(com_df.lazy(), "ts_code", "ts_code")
        .sort(
            ["trend_degree"],
            SortMultipleOptions::new().with_order_descending(true),
        )
        .select([
            col("com_name"),
            col("ts_code"),
            col("up"),
            col("trend_degree")
                .round(2, RoundMode::HalfAwayFromZero)
                .cast(DataType::String)
                + lit("%"),
            col("pct_chg")
                .list()
                .eval(col("").round(2, RoundMode::HalfAwayFromZero), true)
                .list()
                .eval(col("").cast(DataType::String) + lit("%"), true)
                .list()
                .head(lit(11))
                .list()
                .join(lit(","), true)
                .alias("pct_chg_list"),
        ])
        .collect()?;

    let map = get_col_idx!(&frame);
    let com_name_idx = map["com_name"];
    let ts_code_idx = map["ts_code"];

    let diff_idx = map["trend_degree"];
    let pct_chg_list_idx = map["pct_chg_list"];

    let head5 = frame.head(Some(5));

    let mut top5 = String::new();
    for row_idx in 0..head5.height() {
        let row = head5.get_row(row_idx)?;
        let com_name = get_str_col!(row, com_name_idx);
        let com_code = get_str_col!(row, ts_code_idx);
        let diff = get_str_col!(row, diff_idx);
        let pct_chg_list = get_str_col!(row, pct_chg_list_idx);
        let rank = row_idx + 1;
        let msg = format!(
            "上市公司：{com_name}（股票代码{com_code}）5日上行趋势排名第{rank}，趋势度为{diff}，最近11天涨跌情况{pct_chg_list}。\n\n"
        );
        top5.push_str(&msg);
    }

    tracing::info!("\n{top5}");

    let tail5 = frame.tail(Some(5)).reverse();

    let mut last5 = String::new();
    for row_idx in 0..tail5.height() {
        let row = tail5.get_row(row_idx)?;
        let com_name = get_str_col!(row, com_name_idx);
        let com_code = get_str_col!(row, ts_code_idx);
        let diff = get_str_col!(row, diff_idx);
        let pct_chg_list = get_str_col!(row, pct_chg_list_idx);
        let rank = row_idx + 1;
        let msg = format!(
            "上市公司：{com_name}（股票代码{com_code}）5日下行趋势排名第{rank}，趋势度为{diff}，最近11天涨跌情况{pct_chg_list}。\n\n"
        );
        last5.push_str(&msg);
    }

    tracing::info!("\n{last5}");
    
    Ok(vec![top5, last5])
}

fn calc_5d_10d(col: Column) -> PolarsResult<Option<Column>> {
    let pct_chg_list = col.f64()?;
    let emp: Vec<f64> = vec![];
    if pct_chg_list.len() < 11 {
        return Ok(Some(Series::new("5d_10d_lines".into(), emp).into_column()));
    }
    //倒数第二个交易日往前推5个交易日的平均涨跌幅度
    let sum_1: f64 = pct_chg_list.iter().skip(1).filter_map(|v| v).take(5).sum();
    let avg_1 = sum_1 / 5f64;

    //倒数第一个交易日往前推5个交易日的平均涨跌幅度
    let sum_2: f64 = pct_chg_list.iter().filter_map(|v| v).take(5).sum();
    let avg_2 = sum_2 / 5f64;

    //倒数第二个交易日往前推10个交易日的平均涨跌幅度
    let sum_3: f64 = pct_chg_list.iter().skip(1).filter_map(|v| v).take(10).sum();
    let avg_3 = sum_3 / 10f64;

    //倒数第一个交易日往前推10个交易日的平均涨跌幅度
    let sum_4: f64 = pct_chg_list.iter().filter_map(|v| v).take(10).sum();
    let avg_4 = sum_4 / 10f64;

    let vals = vec![avg_1, avg_2, avg_3, avg_4];

    Ok(Some(Series::new("5d_10d_lines".into(), vals).into_column()))
}
